import numpy as np
import matplotlib.pyplot as plt

DNA_SIZE = 10
POP_SIZE = 100          # population size
CROSS_RATE = 0.8        # mating prob,(DNA crossover)
MUTATION_RATE = 0.003   # mutation prob
N_GENERATIONS = 200
X_BOUND = [0, 5]        # x upper and lower bonds


def f(x):
    return np.sin(10 * x) * x + np.cos(2*x) * x


def get_fitness(pred):
    """
    find non-zero fitness for selection
    :param pred:
    :return:
    """
    return pred + 1e-3 - np.min(pred)


def translate_dna(pop):
    """
    convert binary dna to decimal and normalize it to [0,5]
    :param pop: population
    """
    return pop.dot(2 ** np.arange(DNA_SIZE)[::-1]/(2 ** DNA_SIZE -1) * X_BOUND[1])


def select(pop, fitness):
    idx = np.random.choice(np.arange(POP_SIZE), size=POP_SIZE,
                           replace=True, p=fitness/fitness.sum())
    return pop[idx]


def crossover(parent, pop):
    if np.random.rand() < CROSS_RATE:
        i_ = np.random.randint(0, POP_SIZE, size=1)
        cross_point = np.random.randint(0, 2, size=DNA_SIZE).astype(np.bool)
        parent[cross_point] = pop[i_, cross_point]

    return parent


def mutate(child):
    for point in range(DNA_SIZE):
        if np.random.rand() < MUTATION_RATE:
            child[point] = 1 if child[point] == 0 else 0
    return child


if __name__ == '__main__':
    # initialize population DNA
    pop = np.random.randint(0, 2, (1, DNA_SIZE)).repeat(POP_SIZE, axis=0)
    print(pop)
    plt.ion()
    x = np.linspace(*X_BOUND, 200)
    plt.plot(x, f(x))
    for _ in range(N_GENERATIONS):
        # compute values of DNA(transferred to int)
        f_values = f(translate_dna(pop))

        if 'sca' in globals():
            sca.remove()
        sca = plt.scatter(translate_dna(pop), f_values, s=200, lw=0,
                          c='red', alpha=0.5)
        plt.pause(0.05)

        # GA part (evolution)
        fitness = get_fitness(f_values)
        most_fit_dna = pop[np.argmax(fitness), :]
        x_fit = translate_dna(most_fit_dna)
        print("Most fitted DNA: ", most_fit_dna, x_fit)
        pop = select(pop, fitness)
        pop_copy = pop.copy()

        for parent in pop:
            child = crossover(parent, pop_copy)
            child = mutate(child)
            parent[:] = child  # replace parent with child

    plt.ioff()
    plt.show()
